package edu.hit.irlab.coref;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintStream;
import java.io.Writer;
import java.util.ArrayList;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.dom4j.Document;
import org.dom4j.DocumentException;

import edu.hit.irlab.coref.mentionfinder.DefaultCorpus_Mentionfinder;
import edu.hit.irlab.coref.mentionfinder.Mention;
import edu.hit.irlab.coref.mentionfinder.MentionFinder;
import edu.hit.irlab.coref.musicdomain.mentionfinder.CorpusMentionFinder;
import edu.hit.irlab.coref.trainer.CorfTrainer;
import edu.hit.irlab.coref.trainer.SoonTrainer;

import nlpeap.io.XMl;
import nlpeap.ml.InstanceWriter;
import nlpeap.ml.OfflineClassifier;
import nlpeap.ml.maxent.ClassifierSinkBinary;
import nlpeap.ml.maxent.ParameterEstimatorBinary;
import nlpeap.ml.svm.SVMLightClassifier;
import nlpeap.ml.svm.SVMLightInstanceWriter;
import nlpeap.ml.weka.WEKAInstanceClassifier;
import nlpeap.ml.weka.WEKAInstanceWriter;
import nlpeap.toolrepository.LoadException;

/**
 * Train model use classifier
 * @author chuter
 *
 */
public class CRTrainModel {
	public static Log _log = LogFactory.getLog(CRTrainModel.class.getName());
	private static final CRConfigProperties _prop = CRConfigProperties.getInstance();
	private static final String _classifier_type = _prop.getClassifierType();
	private static final String _modelName = _prop.getModelName();
	private static final String _train_data = _prop.getTrainingData();
	
	private InstanceWriter _ins_writer;
	private CorfTrainer _trainer;
	private final MentionFinder<Document> _mFinder;
	private final ArrayList<PairFeatureExtractor> _fes;
	
	public CRTrainModel(
			ArrayList<PairFeatureExtractor> fes,
			MentionFinder<Document> mention_finder)
	{ 
		_fes = fes;
		_mFinder = mention_finder; 
	}
		
	public void setTrainer(CorfTrainer trainer) { _trainer = trainer; }
	public void setInstanceWriter(InstanceWriter ins_writer)
	{ this._ins_writer = ins_writer; }
	
		
	public static void train() throws IOException
	{
		if (_classifier_type.equalsIgnoreCase(OfflineClassifier.WEKA)) {
			if (_modelName != null)
				WEKAInstanceClassifier.setModleName(_modelName);
			WEKAInstanceClassifier.runLearner(
					_train_data, 
					_prop.getWekaClassifierName(),
					_prop.getClassifierOption()
				);
		} else if (_classifier_type.equalsIgnoreCase(OfflineClassifier.SVM)) {
            SVMLightClassifier.runLearner(_modelName, _prop.getClassifierOption());
		} else if (_classifier_type.equalsIgnoreCase(OfflineClassifier.Maxent)) {
			try {
				ParameterEstimatorBinary.do_estimation(new File(_prop.getModelDir(), _modelName).getAbsolutePath());
			} catch (ClassNotFoundException e) {
                e.printStackTrace();
                throw new RuntimeException("Cannot do parameter estimation",e);
			}
		} else {
			_log.fatal("no support for this classifier: "+_classifier_type);
			throw new RuntimeException("There is no such classifier type: "+_classifier_type);
		}
	}
	
	public void createTrainingDataPerDoc(String corpus_file) throws IOException
	{
		Document dom;
		try {
			dom = XMl.loadXmlFile(corpus_file);
		} catch (DocumentException e) {
			_log.fatal(String.format("error in loading the file: %s, please"+
					"check whether the file exists", corpus_file));
			return;
		}
		
		Mention[] mentions = _mFinder.findmentions(dom);
		_trainer.train(mentions);
	}
	
	public void createTrainingData() throws IOException
	{
		BufferedWriter _trainData_writer = new BufferedWriter( new
				FileWriter( new File(_train_data) ));
		
		if (_classifier_type.equalsIgnoreCase(OfflineClassifier.WEKA)) {
			_ins_writer = new WEKAInstanceWriter(_trainData_writer); 
		} else if (_classifier_type.equalsIgnoreCase(OfflineClassifier.SVM)) {
			_ins_writer = new SVMLightInstanceWriter(
					new FileWriter(new File(_prop.getModelDir(), _modelName+".data")),
					new File(_prop.getModelDir(), _modelName+".alphabet"));
		} else if (_classifier_type.equalsIgnoreCase(OfflineClassifier.Maxent)) {
			System.out.println(new File(_prop.getModelDir(), _modelName).getAbsolutePath());
			_ins_writer = new ClassifierSinkBinary(new File(_prop.getModelDir(), _modelName).getAbsolutePath(),
					ClassifierSinkBinary.comboFromString(_prop.getClassifierOption()));
		} else {
			_log.fatal("no support for this classifier: "+_classifier_type);
			//if the classifier type is unknown and the "_ins_writer"
			//hasn't been set either, then throw "RuntimeException"
			if (_ins_writer == null)
				throw new RuntimeException("There is no such classifier type: "+_classifier_type);
		}
		
		_trainer = new SoonTrainer(_fes, _ins_writer);
		
		String corpus_path = _prop.getTrainCorpusPath();
		_log.info(String.format("training model on the corpus [%s]", corpus_path));
		
		File corpusDir = new File(corpus_path);
		for (File file : corpusDir.listFiles()) {
			String filePath = file.getAbsolutePath();
			if (!filePath.endsWith(".xml"))
				continue;
			_log.info(String.format("the trainer running at the file: %s", filePath));
			createTrainingDataPerDoc(filePath);
		}
		
		close();
	}
	
	private void close()
	{
		if (_ins_writer != null) {
			try {
				_ins_writer.close();
			} catch (IOException e) {
				_log.fatal("Can't close the instance writer");
			}
		}
	}
	
	//test program
	public static void main(String[] argv) throws IOException, 
		ClassNotFoundException, InstantiationException,
		IllegalAccessException, LoadException
	{
//		PrintStream out = new PrintStream(new FileOutputStream(new File("."+File.separator+"CRTrainModel.log")));
//		System.setOut(out);
//		System.setErr(out);
		
//		DefaultCorpus_Mentionfinder corp_mFinder = new DefaultCorpus_Mentionfinder();
		CorpusMentionFinder corp_mFinder = new CorpusMentionFinder();
		CRConfigProperties prop = CRConfigProperties.getInstance();
		
		ArrayList<PairFeatureExtractor> fes = prop.getPairFEs();
//		ArrayList<PairFeatureExtractor> fes = new ArrayList<PairFeatureExtractor>();
//		fes.add(new WnZhang());
		
//		Writer insWriter = prop.getTrainInsWriter();
		CRTrainModel ins = new CRTrainModel(fes, corp_mFinder);
		ins.createTrainingData();
		CRTrainModel.train();
	}
}
